#!/usr/bin/env python3
"""
Segmentation point cloud generation utilities for the data processing pipeline.

This module provides a function to generate point clouds with segmentation coloring
from RGB and depth images using camera intrinsics and extrinsics.

Usage:
    This module provides segPcd function for applying segmentation coloring to point cloud data
    within the unified data format.
    
    Mode 2: Use existing point cloud + segmentation images (default behavior)
    Mode 3: Create segmented point cloud from depth images + segmentation images
"""
import numpy as np
import cv2
import argparse
import os
from pathlib import Path
from typing import Optional, Union, Tuple, List
import logging
import warnings
warnings.filterwarnings('ignore')

from unproject import create_pointcloud_from_depth_and_colors, TARGET_CAMERAS

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


def segPcd(unifiedData: dict[str, np.ndarray], 
           depth_scale: float = 0.001) -> dict[str, np.ndarray]:
    """
    Create segmented point cloud from depth images and segmentation images.
    Filters out black pixels from segmentation image (sky/background) by default.
    
    Args:
        unifiedData: dict containing single frame data with keys:
                    ["rgb":np.ndarray,"depth":np.ndarray,"pointcloud":np.ndarray,
                     "segpointcloud":np.ndarray,"state":np.ndarray,"action":np.ndarray,
                     "endpose":np.ndarray,"cameras":dict,"gaussian":dict]
        depth_scale: Scale factor for depth values (default 0.001 for mm to meters)
                    Common values: 0.001 for millimeters to meters, 1.0 for meters
        min_depth: Minimum depth threshold in meters (default 0.1m) - only to filter very close objects
        include_background: If True, include black segmentation pixels (sky/background) (default False)
        
    Returns:
        Dictionary containing original data with updated segpointcloud field
    """
    return create_pointcloud_from_depth_and_colors(
        unifiedData=unifiedData,
        color_key='segmentation',
        output_key='segpointcloud',
        depth_scale=depth_scale,
        min_depth=0.1
    )


def test_pcd_seg():
    """
    Test function for point cloud segmentation.
    """
    # Create dummy test data
    np.random.seed(42)
    
    print("Testing Segmented Point Cloud Generation")
    
    # Create dummy camera data
    cameras_data = {}
    for camera_name in TARGET_CAMERAS:
        # Create dummy segmentation image with black background (sky) and colored objects
        seg_image = np.zeros((240, 320, 3), dtype=np.uint8)  # Start with black background
        # Add some colored objects in the middle
        seg_image[60:180, 80:240] = np.random.randint(50, 255, (120, 160, 3), dtype=np.uint8)  # Colored objects
        
        # Create dummy depth image
        depth_image = np.random.uniform(0.5, 5.0, (240, 320)).astype(np.float32)
        
        # Create dummy camera matrices
        intrinsic_matrix = np.array([
            [320, 0, 160],
            [0, 320, 120], 
            [0, 0, 1]
        ], dtype=np.float32)
        
        # Create dummy extrinsic matrix (identity for simplicity)
        extrinsic_matrix = np.array([
            [1, 0, 0, 0],
            [0, 1, 0, 0], 
            [0, 0, 1, 2]  # 2 meter offset in z
        ], dtype=np.float32)
        
        cameras_data[camera_name] = {
            'segmentation': seg_image,
            'depth': depth_image,
            'intrinsic_cv': intrinsic_matrix,
            'extrinsic_cv': extrinsic_matrix
        }
    
    # Create unified data
    unified_data = {
        'pointcloud': None,
        'segpointcloud': None,
        'cameras': cameras_data,
        'rgb': None,
        'depth': None,
        'state': np.array([]),
        'action': np.array([]),
        'endpose': np.array([]),
        'gaussian': None
    }
    
    # Test without background (filters out black segmentation pixels)
    result_no_bg = segPcd(unified_data.copy())
    print(f"Segmented PCD (exclude black seg pixels) shape: {result_no_bg['segpointcloud'].shape}")
    return result_no_bg


if __name__ == "__main__":
    test_pcd_seg() 